import baostock as bs
import pandas as pd
import os
from datetime import datetime

ADJUST = {'qfq': '1', 'hfq': '2', 'bfq': '3'}


class BaoStockFeed:
    """
    BaoStockFeed
    """
    params = ('date', 'code', 'open', 'high', 'low', 'close', 'volume')

    def __init__(self,
                 saveToCsv=True,
                 save_dir=os.path.join(os.path.dirname(__file__), '..',
                                       'datas')):
        self.saveToCsv = saveToCsv
        self.save_dir = os.path.abspath(save_dir)

        # 初始化连接
        lg = bs.login()
        if lg.error_code != '0':
            print('login baostock server respond error_code:' + lg.error_code)
            print('login baostock server respond  error_msg:' + lg.error_msg)

    def __del__(self):
        bs.logout()

    def get_quote(self,
                  symbol,
                  name,
                  freq='d',
                  start_date='2020-01-01',
                  end_date=None,
                  adjust='qfq'):
        if not end_date:
            end_date = datetime.now().strftime('%Y-%m-%d')

        rs = bs.query_history_k_data_plus(
            symbol,
            ','.join(self.params),  # 所需字段
            start_date=start_date,
            end_date=end_date,
            frequency=freq,  # 周期：d=日k线，w=周k线，m=月k线
            adjustflag=ADJUST[adjust]  # 复权类型：3=不复权，2=后复权，1=前复权
        )

        if rs.error_code != '0':
            print(rs.error_msg)
            return

        # 处理数据
        data_list = []
        while (rs.error_code == '0') & rs.next():
            data_list.append(rs.get_row_data())

        if not data_list:
            print(f'{symbol} {name} {freq} quote data is empty')
            return

        # 转换为DataFrame
        df = pd.DataFrame(data_list, columns=rs.fields)
        df['date'] = pd.to_datetime(df['date'])  # 转换日期格式
        print(df.head())

        if self.saveToCsv:
            self.save(df, f'{symbol}_{name}_{freq}.csv')

    def save(self, df, filename):
        if not os.path.exists(self.save_dir):
            os.mkdir(self.save_dir)

        filename = os.path.join(self.save_dir, filename)
        print(f'save quote data to  {filename}')
        df.to_csv(filename, index=False)


symbols = {'510050.SH': '华夏上证50ETF', '159915.SZ': '易方达创业板ETF'}

if __name__ == '__main__':
    feed = BaoStockFeed()
    for code, name in symbols.items():
        feed.get_quote(code, name)
